package org.virtualdb.socket.proxy;

import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import com.nice.common.memory.BufferPool;

public class AioSockectProxy {

    private int localPort;

    private int serverPort;

    private String serverHost;

    private ExecutorService executor;

    private BufferPool pool = BufferPool.getInstance();

    public AioSockectProxy(int localPort, String serverHost, int serverPort, int bossThreadSize) {
        this.localPort = localPort;
        this.serverPort = serverPort;
        this.serverHost = serverHost;
        this.executor = Executors.newFixedThreadPool(bossThreadSize + 1);
    }

    public AioSockectProxy(int localPort, String serverHost, int serverPort) {
        this(localPort, serverHost, serverPort, Runtime.getRuntime().availableProcessors() + 1);
    }

    public AioSockectProxy(String serverHost, int serverPort) {
        this(1025, serverHost, serverPort, Runtime.getRuntime().availableProcessors() + 1);
    }

    public void start() {
        executor.execute(new Runnable() {
            @Override
            public void run() {
                startListenLocalPort();
            }
        });
    }

    private void startListenLocalPort() {
        try {
            InetSocketAddress local = new InetSocketAddress(localPort);
            AsynchronousChannelGroup group = AsynchronousChannelGroup.withThreadPool(executor);
            AsynchronousServerSocketChannel listener = AsynchronousServerSocketChannel.open(group);
            Future<AsynchronousSocketChannel> accept = listener.bind(local).accept();
            System.out.println(String.format("Proxy [%s:%s] on:[%s]", serverHost, serverPort, localPort));
            doConnectServer(listener, accept.get());
        } catch (Exception e) {
            doException(e, null);
        }
    }

    private void doConnectServer(AsynchronousServerSocketChannel listener, AsynchronousSocketChannel channel) {
        try {
            listener.accept(null, new LocalAccecptHandler(listener));
            InetSocketAddress remote = new InetSocketAddress(serverHost, serverPort);
            AsynchronousSocketChannel serverChannel = AsynchronousSocketChannel.open();
            serverChannel.connect(remote, channel, new ServerConnectHandler(serverChannel));
        } catch (Exception e) {
            doException(e, null);
        }
    }

    protected void ansyPrintData(final SocketAddress from, final SocketAddress to, final ByteBuffer bufReadOnly,
            final int size) {
        executor.execute(new Runnable() {

            @Override
            public void run() {
                int line = 0;
                byte[] array = bufReadOnly.array();
                StringBuilder log = new StringBuilder();
                log.append("[From:").append(from).append("=>To:");
                log.append(to).append(" Size:").append(size).append("]\n");
                log.append("        ").append("| 0  1  2  3  4  5  6  7  8  9  a  b  c  d  e  f |");
                for (int i = 0; i < size; i++) {
                    if (i > 0 && i % 16 == 0) {
                        log.append(String.format("\n|%06x0|", line++));
                        int j = i - 16;
                        while (j < i) {
                            log.append(String.format("%02x ", array[j++]));
                        }
                        log.append("|").append(new String(array, i - 16, 16));
                    }
                }
                System.out.println(log);
            }
        });
    }

    private class BaseReadHandler implements CompletionHandler<Integer, ByteBuffer> {

        protected AsynchronousSocketChannel write;
        protected AsynchronousSocketChannel read;

        public BaseReadHandler(AsynchronousSocketChannel write, AsynchronousSocketChannel read) {
            this.write = write;
            this.read = read;
        }

        @Override
        public void completed(Integer i, ByteBuffer buf) {
            try {
                if (i == -1) {
                    close(read, write);
                    pool.recovery(buf);
                    buf = null;
                    return;
                } else if (i == 0)
                    return;
                buf.flip();
                // ansyPrintData(read.getRemoteAddress(), write.getRemoteAddress(), buf, buf.remaining());
                write.write(buf, buf, new CompletionHandler<Integer, ByteBuffer>() {

                    @Override
                    public void completed(Integer writeSize, ByteBuffer buf) {
                        buf.clear();
                        if (read.isOpen())
                            read.read(buf, buf, BaseReadHandler.this);
                    }

                    @Override
                    public void failed(Throwable e, ByteBuffer buf) {
                        doException(e, buf);
                    }
                });

            } catch (Exception e) {
                doException(e, null);
            }
        }

        @Override
        public void failed(Throwable e, ByteBuffer attachment) {
            close(read, write);
            doException(e, attachment);
        }

        private void close(Closeable... close) {
            if (close == null)
                return;
            for (Closeable closeable : close) {
                try {
                    if (closeable != null)
                        closeable.close();
                } catch (IOException e) {
                    doException(e, closeable);
                }
            }
        }
    }

    private class LocalAccecptHandler implements CompletionHandler<AsynchronousSocketChannel, Void> {

        private AsynchronousServerSocketChannel listener;

        public LocalAccecptHandler(AsynchronousServerSocketChannel listener) {
            this.listener = listener;
        }

        public void completed(AsynchronousSocketChannel localChannel, Void attachment) {
            doConnectServer(listener, localChannel);
        }

        public void failed(Throwable e, Void attachment) {
            doException(e, attachment);
        }

    }

    private class ServerConnectHandler implements CompletionHandler<Void, AsynchronousSocketChannel> {

        private AsynchronousSocketChannel serverChannel;

        public ServerConnectHandler(AsynchronousSocketChannel serverChannel) {
            this.serverChannel = serverChannel;
        }

        public void completed(Void result, AsynchronousSocketChannel localChannel) {
            System.out.println("connected to:" + serverChannel);
            ByteBuffer localBuf = ByteBuffer.allocate(1024);
            ByteBuffer remoteBuf = ByteBuffer.allocate(1024);
            localChannel.read(localBuf, localBuf, new BaseReadHandler(serverChannel, localChannel));
            serverChannel.read(remoteBuf, remoteBuf, new BaseReadHandler(localChannel, serverChannel));
        }

        public void failed(Throwable e, AsynchronousSocketChannel channel) {
            doException(e, channel);
        }

    }

    private void doException(Throwable e, Object channel) {
        e.printStackTrace();
    }

    public static void main(String[] args) {
        if (args.length < 3) {
            System.out.println("Usage: java -jar jsock-proxy.jar  <local port> <remote host> <remote port>");
            return;
        }
        String remoteHost = args[1];
        int localPort = Integer.parseInt(args[0]);
        int remotePort = Integer.parseInt(args[2]);
        if (args.length == 4 && "log".equals(args[3])) {
            System.setProperty("do_log", "true");
        }
        AioSockectProxy monitor = new AioSockectProxy(localPort, remoteHost, remotePort);
        monitor.start();
    }
}
